import base64
import json
from flask import Flask, render_template, request, jsonify, send_from_directory, Response
import os
from openai import OpenAI
from werkzeug.utils import secure_filename

client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

def encode_image(image_path: str) -> str:
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# 初始化 Flask 应用
app = Flask(__name__)

# 上传文件配置
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 检查文件扩展名是否合法
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 生成图片描述
class ResponseGeneratorFromPic:
    def __init__(self, image_path: str):
        self.image_path = image_path
        self.full_response = ""

    def generate_from_pic(self):
        base64_image = encode_image(self.image_path)
        contents = []

        try:
            completion = client.chat.completions.create(
                model="qwen-vl-max",
                messages=[
                    {"role": "user", "content": [
                        {"type": "text", "text": "用图片中识别出来的角色，提供一段system角色的提示词。直接输出提示词即可。"
                         },
                        {"type": "image_url",
                         "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                    ]}
                ],
                stream=True,
            )
            for chunk in completion:
                if hasattr(chunk.choices[0].delta, 'content'):
                    content = chunk.choices[0].delta.content
                    if content:
                        contents.append(content)
                        yield f"data: {content}\n\n"

            yield "data: [DONE]\n\n"

            self.full_response = ''.join(contents)

        except Exception as e:
            yield f"data: Error: {str(e)}\n\n"

class ChatBot:
    def __init__(self, system_desc: str, user_message: str):
        self.system_desc = system_desc
        self.user_message = user_message

    def generate_chat(self):
        try:
            completion = client.chat.completions.create(
                model="qwen-max",
                messages=[
                    {'role': 'system', 'content': self.system_desc},
                    {'role': 'user', 'content': self.user_message}],
                stream=True
                )

            print(self.system_desc)

            for chunk in completion:
                if hasattr(chunk.choices[0].delta, 'content'):
                    content = chunk.choices[0].delta.content
                    if content:
                        yield f"data: {content}\n\n"
            yield "data: [DONE]\n\n"

        except Exception as e:
            yield f"data: Error: {str(e)}\n\n"


response_from_pic = None

# 图片上传路由
@app.route('/upload', methods=['POST'])
def upload_image():

    global response_from_pic

    if 'image' not in request.files:
        return jsonify({'error': '没有检测到图片文件'}), 400

    file = request.files['image']

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        response_from_pic = ResponseGeneratorFromPic(filepath)

        return Response(
            response=response_from_pic.generate_from_pic(),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'X-Accel-Buffering': 'no',
                'Connection': 'keep-alive'
            }
        )
    else:
        return jsonify({'error': '文件格式不支持'}), 400


def generate_chat_response(system_desc, user_message):
    chatbot = ChatBot(system_desc, user_message)
    return chatbot.generate_chat()

# 用户消息路由
@app.route('/chat', methods=['POST'])
def chat_text():
    data = request.get_json()
    user_message = data.get('message', '')

    if user_message:
        print(response_from_pic.full_response)
        response = generate_chat_response(response_from_pic.full_response, user_message)

        return Response(
            response=response,
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'X-Accel-Buffering': 'no',
                'Connection': 'keep-alive'
            }
        )

# 主页路由
@app.route('/')
def index():
    return send_from_directory('static', 'index.html')


# 启动 Flask 应用
if __name__ == '__main__':
    app.run(debug=True)
